import gensim
import numpy as np
import logging

logger = logging.getLogger()


class Word2VecEval(object):
    def __init__(self, model_file):
        self.model = gensim.models.Word2Vec.load(model_file)

    def eval(self, words, topn):
        data = []
        for word in words:
            try:
                out = [ws for ws in self.model.most_similar(word, topn=topn)]
                data.append({"word": word,
                             "similar_words": [ws[0] for ws in out],
                             "similar_words_weight": [ws[1] for ws in out]})

            except Exception as e:
                logger.error('error: %s' % e, exc_info=True)
                data.append({"word": word, "similar_words": [], "similar_words_weight": []})
        return data

    def get_vectors(self, words, agg='sum'):
        vectors = []
        for w in words:
            try:
                vec = self.model[w]
                vectors.append(vec)
            except Exception as e:
                print('word2vec: ', str(e))

        out = np.array(vectors).sum(axis=0)
        if agg != 'sum':
            out = out / len(out)

        return out.tolist()
